﻿using System;
using System.Collections;
using System.Collections.Generic;
using NRank.Algorithms.Strategies;
using NRank.Common;

namespace NRank.Algorithms
{
	/// <summary>
	/// General rankjoin algorithm. Joins any two SORTED collections over joinCondition.
	/// </summary>
	/// <typeparam name="T1">Type of left items</typeparam>
	/// <typeparam name="T2">Type of right items</typeparam>
	/// <typeparam name="T">Type of output items</typeparam>
	public class RankJoin<T1, T2, T> : IEnumerable<IRanked<T>>, IDepthAwareContainer
	{
		/// <summary>
		/// Join function
		/// </summary>
		private readonly Func<T1, T2, T> joinFunc;

		/// <summary>
		/// Pulling strategy
		/// </summary>
		private readonly PullingStrategyTypes pullingStrategy;

		/// <summary>
		/// Aggregate function
		/// </summary>
		private readonly Func<double, double, double> aggregateFunction;

		/// <summary>
		/// Join condition
		/// </summary>
		private readonly Func<T1, T2, bool> joinCondition;
		
		/// <summary>
		/// Left source
		/// </summary>
		//private readonly IEnumerable<IRanked<T1>> leftSource;
		private readonly DepthAware<IRanked<T1>> leftSource;

		/// <summary>
		/// Right source
		/// </summary>
		//private readonly IEnumerable<IRanked<T2>> rightSource;
		private readonly DepthAware<IRanked<T2>> rightSource;

		/// <summary>
		/// Creates instance of RankJoin algorithm
		/// </summary>
		/// <param name="leftSource">Left source</param>
		/// <param name="rightSource">Right source</param>
		/// <param name="rankFunc1">Ranking of left items</param>
		/// <param name="rankFunc2">Ranking of right items</param>
		/// <param name="joinFunc">Join function</param>
		/// <param name="pullingStrategy">Pulling strategy</param>
		/// <param name="aggregateFunc">Aggregate function</param>
		/// <param name="joinCondition">Optional join condition</param>
		public RankJoin(
			IEnumerable<T1> leftSource,
			IEnumerable<T2> rightSource,
			Func<T1, double> rankFunc1,
			Func<T2, double> rankFunc2,
			Func<T1, T2, T> joinFunc,
			PullingStrategyTypes pullingStrategy = PullingStrategyTypes.Alternating,
			Func<double, double, double> aggregateFunc = null,
			Func<T1, T2, bool> joinCondition = null)
		{
			if (leftSource == null) throw new ArgumentNullException("leftSource");
			if (rightSource == null) throw new ArgumentNullException("rightSource");
			if (rankFunc1 == null) throw new ArgumentNullException("rankFunc1");
			if (rankFunc2 == null) throw new ArgumentNullException("rankFunc2");
			if (joinFunc == null) throw new ArgumentNullException("joinFunc");

			this.joinFunc = joinFunc;
			this.pullingStrategy = pullingStrategy;
			this.aggregateFunction = aggregateFunc ?? Defaults.DefaultAggregateFunction;
			this.joinCondition = joinCondition;

			this.leftSource = new DepthAware<IRanked<T1>>(leftSource.AsRanked(rankFunc1));
			this.rightSource = new DepthAware<IRanked<T2>>(rightSource.AsRanked(rankFunc2));

		}

		/// <summary>
		//  Current threshold 
		/// </summary>
		public double Threshold { get; private set; }

		/// <summary>
		/// Algorithm body
		/// </summary>
		/// <returns>Rank-joined items</returns>
		private IEnumerable<IRanked<T>> GetItems()
		{
			Func<IRanked<T1>, IRanked<T2>, IRanked<T>> joinFuncWithRank = (i1, i2) => joinFunc(i1.Object, i2.Object).AsRanked(aggregateFunction(i1.Rank, i2.Rank));

			var joinedResults = new Heap<IRanked<T>>();

			bool leftHasValue = true;
			bool rightHasValue = true;

			var leftAdapter = new RankJoinSource<T1>(leftSource);
			var rightAdapter = new RankJoinSource<T2>(rightSource);

			var left = leftAdapter.GetEnumerator();
			var right = rightAdapter.GetEnumerator();

			var thresholdCounter = new ThresholdCounter(leftAdapter, rightAdapter, this.aggregateFunction);

			IPullingStrategy pullStrategy = Defaults.GetPullStrategy(thresholdCounter, pullingStrategy);

			while (true)
			{
				if (pullStrategy.CanPullLeft && (leftHasValue = left.MoveNext()))
				{
					// pull row from left
					var leftItem = left.Current;

					// update threshold
					Threshold = thresholdCounter.CountThreshold();

					// generate results
					foreach (var rightSeenItem in rightAdapter.SeenItems)
					{
						if (joinCondition == null || joinCondition(leftItem.Object, rightSeenItem.Object))
						{
							joinedResults.Add(joinFuncWithRank(leftItem, rightSeenItem));
						}
					}

					// return valid results
					while (joinedResults.Count > 0 && joinedResults.Peek().Rank >= Threshold)
					{
						yield return joinedResults.RemoveTop();
					}
				}

				if (pullStrategy.CanPullRight && (rightHasValue = right.MoveNext()))
				{
					// pull row from right
					var rightItem = right.Current;

					// update threshold
					Threshold = thresholdCounter.CountThreshold();

					// join with all seen items
					foreach (var leftSeenItem in leftAdapter.SeenItems)
					{
						if (joinCondition == null || joinCondition(leftSeenItem.Object, rightItem.Object))
						{
							var newItem = joinedResults.Add(joinFuncWithRank(leftSeenItem, rightItem));

							bool metNewItem = false;
							while (joinedResults.Count > 0 && joinedResults.Peek().Rank >= Threshold && !metNewItem)
							{
								metNewItem = ReferenceEquals(joinedResults.Peek(), newItem);
								yield return joinedResults.RemoveTop();
							}
						}
					}

					// return valid results
					while (joinedResults.Count > 0 && joinedResults.Peek().Rank >= Threshold)
					{
						yield return joinedResults.RemoveTop();
					}
				}

				// condition to stop
				if ((!pullStrategy.CanPullLeft || !leftHasValue) && (!pullStrategy.CanPullRight || !rightHasValue)) break;
			}

			// just yield the rest of priority queue
			while (joinedResults.Count > 0)
			{
				yield return joinedResults.RemoveTop();
			}
		}

		/// <summary>
		/// Enumerator implementation
		/// </summary>
		/// <returns>Result enumerator</returns>
		public IEnumerator<IRanked<T>> GetEnumerator()
		{
			return GetItems().GetEnumerator();
		}

		/// <summary>
		/// Non-generic enumerator implementation
		/// </summary>
		/// <returns>Non-generic result enumerator</returns>
		IEnumerator IEnumerable.GetEnumerator()
		{
			return GetEnumerator();
		}

		public int LeftDepth
		{
			get { return leftSource.Depth; }
		}

		public int RightDepth
		{
			get { return rightSource.Depth; }
		}
	}
}